{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "20272cd6",
   "metadata": {},
   "source": [
    "cgw-util-cos-sync"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "2db69e62",
   "metadata": {},
   "source": [
    "Syncs two COS buckets using the [CLAIMED Generic Grid Computing Pattern](https://medium.com/@romeokienzler/the-generic-grid-computing-pattern-transforms-any-sequential-workflow-step-into-a-transient-grid-c7f3ca7459c8) which allows:\n",
    "\n",
    "- processing data in partitions based on a file matching pattern for parallelization\n",
    "- running on transient storage only as both sides of this operator are COS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "99f4a7bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install aiobotocore botocore s3fs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f540fcb6",
   "metadata": {},
   "outputs": [],
   "source": [
    "import logging\n",
    "import os\n",
    "import re\n",
    "import s3fs\n",
    "import sys\n",
    "import glob"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9f38e640",
   "metadata": {},
   "outputs": [],
   "source": [
    "# cos path to get job (files) from (including bucket)\n",
    "cgw_source_path = os.environ.get('cgw_source_path') \n",
    "\n",
    "# cgw_source_access_key_id\n",
    "cgw_source_access_key_id = os.environ.get('cgw_source_access_key_id') \n",
    "\n",
    "# source_secret_access_key\n",
    "cgw_source_secret_access_key = os.environ.get('cgw_source_secret_access_key') \n",
    "\n",
    "# source_endpoint\n",
    "cgw_source_endpoint = os.environ.get('cgw_source_endpoint')\n",
    "\n",
    "# cgw_coordinator_access_key_id\n",
    "cgw_coordinator_access_key_id = os.environ.get('cgw_coordinator_access_key_id') \n",
    "\n",
    "# cgw_coordinator_secret_access_key\n",
    "cgw_coordinator_secret_access_key = os.environ.get('cgw_coordinator_secret_access_key') \n",
    "\n",
    "# cgw_coordinator_endpoint\n",
    "cgw_coordinator_endpoint = os.environ.get('cgw_coordinator_endpoint') \n",
    "\n",
    "# cgw_target_access_key_id\n",
    "cgw_target_access_key_id = os.environ.get('cgw_target_access_key_id') \n",
    "\n",
    "# cgw_target_secret_access_key\n",
    "cgw_target_secret_access_key = os.environ.get('cgw_target_secret_access_key') \n",
    "\n",
    "# cgw_target_endpoint\n",
    "cgw_target_endpoint = os.environ.get('cgw_target_endpoint') \n",
    "\n",
    "# cgw_target_path (including bucket)\n",
    "cgw_target_path = os.environ.get('cgw_target_path') \n",
    "\n",
    "# lock file suffix\n",
    "cgw_lock_file_suffix = os.environ.get('cgw_lock_file_suffix', '.lock')\n",
    "\n",
    "# processed file suffix\n",
    "cgw_processed_file_suffix = os.environ.get('cgw_processed_file_suffix', '.processed')\n",
    "\n",
    "# log level\n",
    "cgw_log_level = os.environ.get('cgw_log_level', 'INFO')\n",
    "\n",
    "# timeout in seconds to remove lock file from struggling job (default 1 hour)\n",
    "cgw_lock_timeout = int(os.environ.get('cgw_lock_timeout', 60*60))\n",
    "\n",
    "# group files which need to be processed together\n",
    "cgw_group_by = os.environ.get('cgw_group_by', None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "87541c44",
   "metadata": {},
   "outputs": [],
   "source": [
    "root = logging.getLogger()\n",
    "root.setLevel(cgw_log_level)\n",
    "\n",
    "handler = logging.StreamHandler(sys.stdout)\n",
    "handler.setLevel(cgw_log_level)\n",
    "formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')\n",
    "handler.setFormatter(formatter)\n",
    "root.addHandler(handler)\n",
    "\n",
    "logging.basicConfig(level=logging.CRITICAL)\n",
    "\n",
    "parameters = list(\n",
    "    map(lambda s: re.sub('$', '\"', s),\n",
    "        map(\n",
    "            lambda s: s.replace('=', '=\"'),\n",
    "            filter(\n",
    "                lambda s: s.find('=') > -1 and bool(re.match(r'[A-Za-z0-9_]*=[.\\/A-Za-z0-9]*', s)),\n",
    "                sys.argv\n",
    "            )\n",
    "    )))\n",
    "\n",
    "for parameter in parameters:\n",
    "    logging.info('Parameter: ' + parameter)\n",
    "    exec(parameter)\n",
    "\n",
    "# parameter postprocessing\n",
    "cgw_source_path = cgw_source_path[1:] if cgw_source_path.startswith('/') else cgw_source_path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5c298b81",
   "metadata": {},
   "outputs": [],
   "source": [
    "input_path = 'input'\n",
    "target_path = 'target'\n",
    "\n",
    "if not os.path.exists(input_path):\n",
    "    os.makedirs(input_path)\n",
    "    \n",
    "if not os.path.exists(target_path):\n",
    "    os.makedirs(target_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b3885a51",
   "metadata": {},
   "outputs": [],
   "source": [
    "def process():\n",
    "    return os.listdir(target_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3a845721",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import s3fs\n",
    "\n",
    "s3source = s3fs.S3FileSystem(\n",
    "    anon=False,\n",
    "    key=cgw_source_access_key_id,\n",
    "    secret=cgw_source_secret_access_key,\n",
    "    client_kwargs={'endpoint_url': cgw_source_endpoint})\n",
    "\n",
    "s3target = s3fs.S3FileSystem(\n",
    "    anon=False,\n",
    "    key=cgw_target_access_key_id,\n",
    "    secret=cgw_target_secret_access_key,\n",
    "    client_kwargs={'endpoint_url': cgw_target_endpoint})\n",
    "\n",
    "s3coordinator = s3fs.S3FileSystem(\n",
    "    anon=False,\n",
    "    key=cgw_coordinator_access_key_id,\n",
    "    secret=cgw_coordinator_secret_access_key,\n",
    "    client_kwargs={'endpoint_url': cgw_coordinator_endpoint})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "72fb08fd",
   "metadata": {},
   "outputs": [],
   "source": [
    "groups = set()\n",
    "all_files = s3source.glob(cgw_source_path)\n",
    "for path_string in all_files:\n",
    "    exec('part = path_string' + cgw_group_by)\n",
    "    logging.info(f'Identified group: {part}')\n",
    "    groups.add(part)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aacd80ef",
   "metadata": {},
   "outputs": [],
   "source": [
    "from datetime import datetime, timedelta\n",
    "import os\n",
    "import random\n",
    "import time\n",
    "\n",
    "delay = random.randint(1, 60)\n",
    "logging.info(f'Staggering start, waiting for {delay} seconds')\n",
    "time.sleep(delay)\n",
    "\n",
    "for batch in groups:\n",
    "\tlogging.info(f'Processing batch: {batch}')\n",
    "\n",
    "\tbatch_fileset = list(filter(lambda file: batch in file, all_files))\n",
    "\tfile = batch_fileset[0]\n",
    "\n",
    "\tlogging.info(f'Processing {file}: ')\n",
    "\n",
    "\tlock_file = file+cgw_lock_file_suffix\n",
    "\tprocessed_file = file+cgw_processed_file_suffix\n",
    "\n",
    "\tdef is_locked():\n",
    "\t\texists =s3coordinator.exists(lock_file)\n",
    "\t\tlogging.info(f'Lock file {lock_file} {exists}')\n",
    "\t\treturn exists\n",
    "\n",
    "\tdef is_processed():\n",
    "\t\texists =s3coordinator.exists(processed_file)\n",
    "\t\tlogging.info(f'Processed file {processed_file} {exists}')\n",
    "\t\treturn exists\n",
    "\n",
    "\tdef unlock():\n",
    "\t\ts3coordinator.rm(lock_file)\n",
    "\n",
    "\tdef lock():\n",
    "\t\tlogging.info(f'Locking {lock_file}')\n",
    "\t\ts3coordinator.touch(lock_file)\n",
    "\n",
    "\tdef is_locked_expired():\n",
    "\t\tif is_locked():\n",
    "\t\t\tlast_modified = s3coordinator.info(lock_file)['LastModified']\n",
    "\t\t\tlogging.info(f'Lock file {lock_file} last modified {last_modified}')\n",
    "\t\t\ttime_threshold = datetime.now(last_modified.tzinfo) - timedelta(seconds=cgw_lock_timeout)\n",
    "\t\t\tlogging.info(f'Lock file {lock_file} time threshold {time_threshold}')\n",
    "\t\t\texpired = last_modified < time_threshold\n",
    "\t\t\tlogging.info(f'Lock file {lock_file} expired {expired}')\n",
    "\t\t\treturn expired\n",
    "\t\telse:\n",
    "\t\t\treturn False\n",
    "\n",
    "\t# remove strugglers\n",
    "\tif is_locked_expired():\n",
    "\t\ts3coordinator.rm(lock_file)\n",
    "\n",
    "\tif is_locked():\n",
    "\t\tcontinue\n",
    "\n",
    "\tif is_processed():\n",
    "\t\tcontinue\n",
    "\t\n",
    "\tlock()\n",
    "\n",
    "\t# downloading input files to local storage\n",
    "\tfor batch_file in batch_fileset:\n",
    "\t\tlogging.info(f'Downloading {batch_file}')\n",
    "\t\t#if os.path.exists(input_path + '/' + batch_file.split('/')[-1]):\n",
    "\t\t#\tcontinue\n",
    "\t\ts3source.get(batch_file,input_path + '/' + batch_file.split('/')[-1])\n",
    "\n",
    "\toutput_files : list = process()\n",
    "\n",
    "\t# removing local input files\n",
    "\tfor batch_file in batch_fileset:\n",
    "\t\tfile_to_remove = batch_file.split('/')[-1]\n",
    "\t\tlogging.info(f'Removing {file_to_remove}')\n",
    "\t\tos.remove(os.path.join(input_path,file_to_remove))\n",
    "\n",
    "\t# uploading results to cos\n",
    "\tfor output_file in output_files:\n",
    "\t\tlogging.info(f'Uploading {output_file}')\n",
    "\t\ts3target.put(os.path.join(target_path,output_file),cgw_target_path + '/' + output_file)\n",
    "\n",
    "\t# removing local target files\n",
    "\tfor output_file in output_files:\n",
    "\t\toutput_file = output_file.split('/')[-1]\n",
    "\t\tlogging.info(f'Removing {output_file}')\n",
    "\t\tos.remove(os.path.join(target_path,output_file))\n",
    "\t\n",
    "\t# marking input files as processed\n",
    "\tlogging.info(f'Marking {processed_file} as processed')\n",
    "\ts3source.touch(processed_file)\n",
    "\t\n",
    "\tunlock()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15972ce2",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
